import torch
import torch.nn as nn
import numpy as np
from prob import *

torch.set_default_dtype(torch.float64)
torch.set_default_tensor_type(torch.DoubleTensor)


def gen_nn():
    act_fun = nn.ReLU()
    layer_input = [nn.Linear(DIM, D_H)]
    layer_output = [act_fun, nn.Linear(D_H, 1)]

    # hidden layer
    module_hidden = [[act_fun, nn.Linear(D_H, D_H)] for _ in range(N_H - 1)]
    layer_hidden = list(np.array(module_hidden).flatten())

    # nn model
    layers = layer_input + layer_hidden + layer_output
    model = nn.Sequential(*layers)

    return model


def serve_for_verify(model):
    w_b_list = []
    for name, p in model.named_parameters():
        w_b_list.append(p.data.cpu().detach().numpy())
    return w_b_list
